import time
import numpy as np
import matplotlib.pyplot as plt
from sklearn import (manifold, datasets, decomposition, ensemble, discriminant_analysis, random_projection)
import torch
import torch.nn.functional as F


def plot_embedding(X, y, title=None):
    x_min, x_max = np.min(X, 0), np.max(X, 0)
    X = (X - x_min) / (x_max - x_min)

    plt.figure()
    # ax = plt.subplot(111)
    for i in range(X.shape[0]):
        plt.text(X[i, 0], X[i, 1], str(y[i]), color=plt.cm.Set1(y[i] / 10.), fontdict={'weight': 'bold', 'size': 9})

    shown_images = np.array([[1., 1.]])  # just something big
    for i in range(X.shape[0]):
        dist = np.sum((X[i] - shown_images) ** 2, 1)
        # if np.min(dist) < 4e-3:
        #     # don't show points that are too close
        #     continue
        if np.max(dist) > 4e-3:
            # don't show points that are too far
            continue
        shown_images = np.r_[shown_images, [X[i]]]
    plt.xticks([]), plt.yticks([])
    plt.title(title)
    time_str = time.strftime("%Y%m%d") + time.strftime("_%H%M%S")
    plt.savefig(time_str + ".png")
    plt.show()


'''
三种方法
'''


def tSNE(X, y, title):
    print("Computing t-SNE embedding")
    # tsne = manifold.TSNE(n_components=2, init='pca', random_state=0)
    tsne = manifold.TSNE()
    X_tsne = tsne.fit_transform(X)
    plot_embedding(X_tsne, y, "tSNE" + title)


def MDS(X, y, title):
    print("Computing MDS embedding")
    clf = manifold.MDS(n_components=2, n_init=1, max_iter=100)
    X_mds = clf.fit_transform(X)
    plot_embedding(X_mds, y, "MDS" + title)


def RandomForest(X, y, title):
    print("Computing Totally Random Trees embedding")
    hasher = ensemble.RandomTreesEmbedding(n_estimators=200, random_state=0, max_depth=5)
    X_transformed = hasher.fit_transform(X)
    pca = decomposition.TruncatedSVD(n_components=2)
    X_reduced = pca.fit_transform(X_transformed)
    plot_embedding(X_reduced, y, "RandomForest" + title)
